import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import tqdm
from models.score_net import ScoreNet
from torch.utils.data import DataLoader
from models.sde_lib import VPSDE, subVPSDE

if __name__ == "__main__":
    
    model = ScoreNet()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.train()

    # load dataset cifar100
    train_dataset = torchvision.datasets.CIFAR100(root="./data", train=True, transform=torchvision.transforms.ToTensor(), download=False)
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    otimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    sde = subVPSDE(beta_min=0.1, beta_max=20, N=1000)
    step = 0
    for epoch in range(1000):
        for idx, (images, _) in tqdm.tqdm(enumerate(train_loader), desc="Training", total=len(train_loader)):
            images = images.to(device)
            #images = torch.mean(images, dim=1, keepdim=True)
            t = torch.rand(images.shape[0], device=device)
            t = t * 999

            z = torch.randn_like(images)
            mean, std = sde.marginal_prob(images, t)
            perturbed_images = images + std[:, None, None, None] * z
            # forward pass
            score = model(perturbed_images, t)
            std = sde.marginal_prob(torch.zeros_like(images), t)[1]
            
            score = -score / std[:, None, None, None]

            losses = torch.square(score * std[:, None, None, None] + z)

            loss = torch.sum(losses)
            
            step += 1
            if step % 100 == 0:
                print(f"Epoch: {epoch}, Step: {step}, Loss: {loss.item()}")
            otimizer.zero_grad()
            loss.backward()
    